import re

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from paddleocr import PaddleOCR

# 初始化 PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang='ch', det_model_dir='./ocr_model/det/ch/ch_PP-OCRv4_det_infer',
                rec_model_dir='./ocr_model/rec/ch/ch_PP-OCRv4_rec_infer',
                cls_model_dir='./ocr_model/cls/ch_ppocr_mobile_v2.0_cls_infer')  # 使用中文模型


def add_mosaic(img, x, y, w, h, mosaic_size=10):
    # 确保选择的区域在图像内
    if x < 0 or x + w > img.shape[1] or y < 0 or y + h > img.shape[0]:
        print("选择的区域超出图像边界，无法添加马赛克")
        return img

    sub_img = img[y:y + h, x:x + w]  # 身份证号区域

    if sub_img.size == 0:
        print("选择的区域为空，无法添加马赛克")
        return img

    sub_img_small = cv2.resize(sub_img, (mosaic_size, mosaic_size), interpolation=cv2.INTER_LINEAR)
    sub_img_mosaic = cv2.resize(sub_img_small, (w, h), interpolation=cv2.INTER_NEAREST)

    img[y:y + h, x:x + w] = sub_img_mosaic
    return img


def apply_mosaic(image, points, mosaic_size=10):
    """
    在指定的矩形区域内应用马赛克效果。

    :param image: 输入图像
    :param points: 四个标记点的坐标 (x, y)，形成一个矩形区域
    :param mosaic_size: 马赛克块的大小
    :return: 处理后的图像
    """
    # 计算矩形区域的边界
    x, y, w, h = cv2.boundingRect(np.array(points, dtype=np.int32))

    # 提取矩形区域
    region = image[y:y + h, x:x + w]

    # 缩小和放大以实现马赛克效果
    small = cv2.resize(region, (mosaic_size, mosaic_size), interpolation=cv2.INTER_LINEAR)
    mosaic = cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)

    # 将马赛克区域放回原图
    image[y:y + h, x:x + w] = mosaic

    return image


def rotate_image(image, angle):
    """ 旋转图片 """
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    return rotated


def detect_orientation(image):
    """ 检测图片方向 """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)

    angles = []
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 计算线段的长度
            length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            # 忽略较短的线段
            if length > 50:
                angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
                angles.append(angle)

    if angles:
        # 计算角度的中位数，减少异常值的影响
        median_angle = np.median(angles)
        # 将角度限制在 [-90, 90] 范围内
        median_angle = (median_angle + 180) % 180 - 90
        return median_angle
    else:
        return 0


def process_image(image_path):
    """ 处理图片，包括方向检测和矫正 """
    image = cv2.imread(image_path)
    angle = detect_orientation(image)

    if angle != 0:
        image = rotate_image(image, -angle)

    return image, angle


def get_card_type(rec_text):
    # 要检查的子字符串列表
    id_key_word = ['公民身份号码']
    license_key_word = ["法定代表人"]
    weight_ticket_key_word = ["过磅单"]
    bank_advice = ["回单"]
    bank_key_word = ["Debit", "DEBIT", "ATM", "银联", "BOC", "CCB", "ABC", "ICBC", "BOCOM", "CMB", "CIB", "CMBC", "HXB",
                     "PAB",
                     "CEB", "CNCB", "SPDB"]

    # 使用any()函数检查列表中是否有元素包含任意一个关键词 exists =
    if any(any(key in item for item in rec_text) for key in id_key_word):
        return {"cardType": "ID"}
    elif any(any(key in item for item in rec_text) for key in license_key_word):
        return {"cardType": "License"}
    elif any(any(key in item for item in rec_text) for key in weight_ticket_key_word):
        return {"cardType": "WeightTicket"}
    elif any(any(key in item for item in rec_text) for key in bank_advice):
        return {"cardType": "BankAdvice"}
    elif any(any(key in item for item in rec_text) for key in bank_key_word):
        return {"cardType": "Bank"}
    else:
        return {"cardType": "Unknown"}


def put_chinese_text(image, text, position, font_path, font_size, color=(255, 255, 255)):
    """
    在图像上绘制中文文本。

    :param image: 输入图像（OpenCV 格式）
    :param text: 要绘制的文本
    :param position: 文本的位置 (x, y)
    :param font_path: 字体文件的路径
    :param font_size: 字体大小
    :param color: 文本颜色，默认为白色 (B, G, R)
    :return: 绘制文本后的图像
    """
    # 将 OpenCV 图像转换为 PIL 图像
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    # 创建绘图对象
    draw = ImageDraw.Draw(pil_image)

    # 加载字体
    font = ImageFont.truetype(font_path, font_size)

    # 绘制文本
    draw.text(position, text, fill=color, font=font)

    # 将 PIL 图像转换回 OpenCV 图像
    cv2_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

    return cv2_image


def draw_boxes(image, boxes, txts, scores, font_path='./fonts/chinese_cht.ttf'):
    # 遍历检测结果并绘制多边形和文本
    for box, txt, score in zip(boxes, txts, scores):
        pts = np.array(box, np.int32)
        pts = pts.reshape((-1, 1, 2))

        # 绘制多边形
        cv2.polylines(image, [pts], isClosed=True, color=(0, 0, 255), thickness=2)

        # 计算文本位置
        text_position = (pts.min(axis=0)[0][0], pts.min(axis=0)[0][1] - 10)  # 放在多边形上方一点

        # 构建完整的文本字符串
        full_text = f"{txt} ({score:.2f})"

        # 绘制文本
        image = put_chinese_text(image, full_text, text_position, font_path, font_size=10, color=(0, 0, 255))

    return image


def draw_mosaic_id(image, boxes, txts, scores, mosaic_size=10):
    id_number_pattern = r'\d{17}[\dX]'
    for box, txt, score in zip(boxes, txts, scores):
        if re.search(id_number_pattern, txt):
            id_number = re.search(id_number_pattern, txt).group()
            mo_image = apply_mosaic(image, box, mosaic_size=mosaic_size)
            return id_number, mo_image


if __name__ == '__main__':
    # 读取身份证图片
    image_path = './image/4.png'
    # corrected_image, detected_angle = process_image(image_path)

    # 打印检测到的角度
    # print(f"Detected Angle: {detected_angle}")

    # 如果检测到的角度为0，手动指定一个角度进行测试
    # if detected_angle == 0:
    #     corrected_image = rotate_image(corrected_image, -90)  # 手动指定一个角度（例如90度）

    image = cv2.imread(image_path)
    # 进行 OCR 识别
    result = ocr.ocr(image, cls=True)
    result = result[0]
    # 打印结果的前几项，查看数据结构
    print("OCR Result Structure:")
    print(result)  # 打印前5个结果

    boxes = [line[0] for line in result if len(line) == 2]
    txts = [line[1][0] for line in result if len(line) == 2]
    scores = [line[1][1] for line in result if len(line) == 2]
    print(boxes)
    print(txts)
    # 绘制边界框和文本
    # box_image = draw_boxes(image, boxes, txts, scores)
    # cv2.imshow("ID Card with Mosaic", box_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # 显示图像  马赛克
    # _, mo_image = draw_mosaic_id(image, boxes, txts, scores, mosaic_size=15)
    # cv2.imshow("ID Card with Mosaic", mo_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # 可视化结果
    # image_pil = Image.fromarray(cv2.cvtColor(corrected_image, cv2.COLOR_BGR2RGB))
    # im_show = draw_ocr(image_pil, boxes, txts, scores)
    # im_show = Image.fromarray(im_show)
    # im_show.save('result.jpg')
    # im_show.show()

    # f分类测试
    card_type = get_card_type(txts)
    print(card_type)
